from tushare_api import TushareClient
import pandas as pd


def get_kline_data(client, ts_codes, trade_date):
    """
    批量获取指定交易日的 K 线数据
    :param client: TushareClient 实例
    :param ts_codes: 逗号分隔的股票代码字符串
    :param trade_date: 交易日期 (YYYYMMDD)
    :return: 返回包含 open, close, high, low 的 DataFrame
    """
    try:
        # 调用 API 获取 K 线数据
        df_kline = client.get_daily_kline(ts_codes, trade_date=trade_date)

        # 如果返回为空，则返回空 DataFrame
        if df_kline.empty:
            return pd.DataFrame(columns=['ts_code', 'open', 'close', 'high', 'low'])

        return df_kline[['ts_code', 'open', 'close', 'high', 'low']]

    except Exception as e:
        print(f"Error fetching data for {trade_date}: {e}")
        return pd.DataFrame(columns=['ts_code', 'open', 'close', 'high', 'low'])


if __name__ == '__main__':
    client = TushareClient()
    df = pd.read_excel("/Users/decheng/Desktop/涨停双响炮次日溢价情况.xlsx")

    # 向量化生成 ts_code
    df['ts_code'] = df['股票代码'].str.replace('\'', '').apply(
        lambda x: x + '.SH' if x.startswith('60') else x + '.SZ'
    )

    # 按交易日分组，收集所有 ts_code
    grouped = df.groupby('终止时间')['ts_code'].apply(lambda x: ','.join(x)).reset_index(name='ts_codes')

    # 批量获取 K 线数据
    kline_results = []
    for _, row in grouped.iterrows():
        trade_date = row['终止时间']
        ts_codes = row['ts_codes']
        df_kline = get_kline_data(client, ts_codes, trade_date)
        df_kline['trade_date'] = trade_date  # 添加交易日列
        kline_results.append(df_kline)

    # 合并所有 K 线数据
    combined_kline = pd.concat(kline_results, ignore_index=True)

    # 合并回原始 DataFrame
    df = df.merge(
        combined_kline,
        left_on=['ts_code', '终止时间'],
        right_on=['ts_code', 'trade_date'],
        how='left'
    )

    # 重命名并填充缺失值
    df.rename(columns={
        'open': '开盘价',
        'close': '收盘价',
        'high': '最高价',
        'low': '最低价'
    }, inplace=True)
    df.fillna(0, inplace=True)

    # 输出结果
    print(df.head())
    df.to_excel("/Users/decheng/Desktop/export1.xlsx", index=False)
